/*******************************************************************************
 * Copyright 2009 DCSpectrometer - http://code.google.com/p/dcspectrometer 
 *  
 * Licensed under the Apache License, Version 2.0 (the "License"); 
 * you may not use this file except in compliance with the License. 
 * You may obtain a copy of the License at 
 *  
 *     http://www.apache.org/licenses/LICENSE-2.0 
 *     
 * Unless required by applicable law or agreed to in writing, software 
 * distributed under the License is distributed on an "AS IS" BASIS, 
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 * See the License for the specific language governing permissions and 
 * limitations under the License. 
 *******************************************************************************/
package com.dcspectrometer.kmeansclustering;

import java.awt.Point;
import java.awt.geom.Point2D;
import java.util.ArrayList;

public class KMeanClustering {
  private Cluster[] _clusters;
  private ArrayList<Point> _points;
  private double _wcss;

  /**
   * Construct k-mean clustering model
   * 
   * @param k
   *          number of means
   * @param points
   *          data points to cluster
   */
  public KMeanClustering(int k, ArrayList<Point> points) {
    _clusters = new Cluster[k];
    for (int i = 0; i < k; i++) {
      _clusters[i] = new Cluster();
    }
    _points = points;

    performClustering();
  }

  /**
   * Recalculate within-cluster sum of squares
   */
  private void recalculateWCSS() {
    _wcss = 0;
    for (int i = 0; i < _clusters.length; i++) {
      _wcss += _clusters[i].getSumOfDistanceSquares();
    }
  }

  /**
   * 
   * @return
   */
  public ArrayList<Cluster> getClusters() {
    ArrayList<Cluster> clusterPoints = new ArrayList<Cluster>();
    for (int i = 0; i < _clusters.length; i++) {
      clusterPoints.add(_clusters[i]);
    }
    return clusterPoints;
  }

  /**
   * Returns number of means
   * 
   * @return
   */
  public int getKValue() {
    return _clusters.length;
  }

  /**
   * Returns number of points
   * 
   * @return
   */
  public int getNumberOfPoints() {
    return _points.size();
  }

  /**
   * Returns within-cluster sum of squares
   * 
   * @return
   */
  public double getWCSS() {
    recalculateWCSS();
    return _wcss;
  }

  /**
   * Returns a particular cluster by index.
   * 
   * @param index
   *          index of the cluster
   * @return
   */
  public Cluster getCluster(int index) {
    return _clusters[index];
  }

  private void pickInitialCentroids() {
    ArrayList<Point> centroidLocations = new ArrayList<Point>();
    ArrayList<Point> otherPoints = new ArrayList<Point>();

    /** Take any point (why not the first one) as a first centroid location. */
    centroidLocations.add(_points.get(0));

    /** Other points are put on the second list */
    for (int i = 1; i < _points.size(); i++) {
      otherPoints.add(_points.get(i));
    }

    centroidLocations = getCentroids(centroidLocations, otherPoints);

    for (int i = 0; i < centroidLocations.size(); i++) {
      // System.out.println("Centroid " + i + " is at " + centroidLocations.get(i).x + ", "
      // + centroidLocations.get(i).y);

      ClusterCentroid nextCentroid = new ClusterCentroid(centroidLocations.get(i));
      _clusters[i].setCentroid(nextCentroid);
      nextCentroid.setCluster(_clusters[i]);
    }
  }

  /**
   * Recursively looks for k points that are "farthest apart". An energy function
   * sum(1/(distance^2)) is used.
   * 
   * @param centroidLocations
   *          centroid location that have been already selected.
   * @param otherPoints
   *          available choices for next centroid.
   * @return k centroids that are "significantly spread out".
   */
  private ArrayList<Point> getCentroids(ArrayList<Point> centroidLocations,
      ArrayList<Point> otherPoints) {
    /** Check if we done. */
    if (centroidLocations.size() >= _clusters.length) {
      return centroidLocations;
    }

    /** Calculate energy to other points. */
    double[] pointEnergy = new double[otherPoints.size()];

    for (int i = 0; i < otherPoints.size(); i++) {
      pointEnergy[i] = 0;
      for (int j = 0; j < centroidLocations.size(); j++) {
        pointEnergy[i] += 1 / otherPoints.get(i).distanceSq(centroidLocations.get(j));
      }
    }

    /** Find point with lowest energy (farthest point) */
    double minEnergy = pointEnergy[0];
    Point minEnergyPoint = otherPoints.get(0);
    for (int i = 1; i < otherPoints.size(); i++) {
      if (pointEnergy[i] < minEnergy) {
        minEnergy = pointEnergy[i];
        minEnergyPoint = otherPoints.get(i);
      }
    }

    otherPoints.remove(minEnergyPoint);
    centroidLocations.add(minEnergyPoint);

    return getCentroids(centroidLocations, otherPoints);
  }

  public void recalculateCentroids() {
    /** Find new location of cluster centroids. */
    int x, y;
    for (int i = 0; i < _clusters.length; i++) {
      x = 0;
      y = 0;
      for (int j = 0; j < _clusters[i].getPointCount(); j++) {
        Point nextPoint = _clusters[i].getPoint(j);
        x += nextPoint.x;
        y += nextPoint.y;
      }
      x /= _clusters[i].getPointCount();
      y /= _clusters[i].getPointCount();

      _clusters[i].setCentroid(new Point2D.Double(x, y));
    }
  }

  /**
   * Assigns points to clusters and returns number of reassignments
   * 
   * @return
   */
  private int clusterPointsAccordingToCentroids(boolean initialAssignment) {
    int numberOfReasignments = 0;

    if (initialAssignment) {
      for (int i = 0; i < _points.size(); i++) {
        Point nextPoint = _points.get(i);
        _clusters[findClosestClusterIndex(nextPoint)].addPoint(nextPoint);
        numberOfReasignments++;
      }
    } else {
      for (int i = 0; i < _clusters.length; i++) {
        for (int j = 0; j < _clusters[i].getPointCount(); j++) {
          Point nextPoint = _clusters[i].getPoint(j);
          int closestClusterIndex = findClosestClusterIndex(nextPoint);

          if (closestClusterIndex != i) {
            _clusters[i].removePoint(nextPoint);
            _clusters[closestClusterIndex].addPoint(nextPoint);
            numberOfReasignments++;
          }
        }
      }
    }

    return numberOfReasignments;
  }

  private int findClosestClusterIndex(Point point) {
    int minIndex = 0;
    double minDistance = point.distanceSq(_clusters[0].getCentroid());

    for (int i = 1; i < _clusters.length; i++) {
      double nextDistance = point.distanceSq(_clusters[i].getCentroid());
      if (minDistance > nextDistance) {
        minIndex = i;
        minDistance = nextDistance;
      }
    }

    return minIndex;
  }

  private void performClustering() {

    /** Find initial centroids */
    pickInitialCentroids();

    /** Perform initial point assignment */
    int numberOfReassignments = clusterPointsAccordingToCentroids(true);

    while (numberOfReassignments > 0) {
      /** Recalculate centroids */
      recalculateCentroids();

      /** Perform initial point assignment */
      numberOfReassignments = clusterPointsAccordingToCentroids(false);
    }
  }
}
